package com.example.test;

import org.junit.BeforeClass;
import org.junit.Test;
import org.openhitls.crypto.jce.HiTls4jProvider;
import javax.crypto.Mac;
import javax.crypto.spec.SecretKeySpec;
import java.nio.charset.StandardCharsets;
import java.security.Security;
import static org.junit.Assert.*;

public class HMACSM3JceTest {
    @BeforeClass
    public static void setUp() {
        Security.addProvider(new HiTls4jProvider());
    }

    @Test
    public void testHMACSM3SingleShot() throws Exception {
        byte[] key = "TestKey123".getBytes(StandardCharsets.UTF_8);
        SecretKeySpec keySpec = new SecretKeySpec(key, "HMACSM3");
        Mac mac = Mac.getInstance("HMACSM3", "HITLS4J");
        mac.init(keySpec);

        String message = "Hello, HMAC-SM3!";
        byte[] macResult = mac.doFinal(message.getBytes(StandardCharsets.UTF_8));

        // Verify MAC length
        assertEquals("MAC length should be 32 bytes", 32, macResult.length);
    }

    @Test
    public void testHMACSM3Incremental() throws Exception {
        byte[] key = "TestKey123".getBytes(StandardCharsets.UTF_8);
        SecretKeySpec keySpec = new SecretKeySpec(key, "HMACSM3");
        Mac mac = Mac.getInstance("HMACSM3", "HITLS4J");
        mac.init(keySpec);

        String part1 = "Hello, ";
        String part2 = "HMAC-SM3!";

        // Incremental update
        mac.update(part1.getBytes(StandardCharsets.UTF_8));
        mac.update(part2.getBytes(StandardCharsets.UTF_8));
        byte[] incrementalMac = mac.doFinal();

        // Compare with single-shot MAC
        mac.reset();
        byte[] singleMac = mac.doFinal((part1 + part2).getBytes(StandardCharsets.UTF_8));

        assertArrayEquals("Incremental and single-shot MACs should match", singleMac, incrementalMac);
    }

    @Test
    public void testEmptyMessage() throws Exception {
        byte[] key = "TestKey123".getBytes(StandardCharsets.UTF_8);
        SecretKeySpec keySpec = new SecretKeySpec(key, "HMACSM3");
        Mac mac = Mac.getInstance("HMACSM3", "HITLS4J");
        mac.init(keySpec);

        byte[] macResult = mac.doFinal(new byte[0]);
        assertEquals("MAC length should be 32 bytes", 32, macResult.length);
    }

    @Test
    public void testReset() throws Exception {
        byte[] key = "TestKey123".getBytes(StandardCharsets.UTF_8);
        SecretKeySpec keySpec = new SecretKeySpec(key, "HMACSM3");
        Mac mac = Mac.getInstance("HMACSM3", "HITLS4J");
        mac.init(keySpec);

        String message = "Test message";
        byte[] mac1 = mac.doFinal(message.getBytes(StandardCharsets.UTF_8));

        // Get another MAC after reset
        mac.reset();
        byte[] mac2 = mac.doFinal(message.getBytes(StandardCharsets.UTF_8));

        assertArrayEquals("MACs should be identical after reset", mac1, mac2);
    }

    @Test
    public void testMultiThread() throws Exception {
        final String message = "Test message";
        final byte[] key = "TestKey123".getBytes(StandardCharsets.UTF_8);
        final SecretKeySpec keySpec = new SecretKeySpec(key, "HMACSM3");

        Mac mac = Mac.getInstance("HMACSM3", "HITLS4J");
        mac.init(keySpec);
        final byte[] expectedMac = mac.doFinal(message.getBytes(StandardCharsets.UTF_8));

        Thread[] threads = new Thread[2];
        final boolean[] results = new boolean[threads.length];

        for (int i = 0; i < threads.length; i++) {
            final int index = i;
            threads[i] = new Thread(() -> {
                try {
                    Mac threadMac = Mac.getInstance("HMACSM3", "HITLS4J");
                    threadMac.init(keySpec);
                    byte[] threadMacResult = threadMac.doFinal(message.getBytes(StandardCharsets.UTF_8));
                    results[index] = java.util.Arrays.equals(expectedMac, threadMacResult);
                } catch (Exception e) {
                    results[index] = false;
                }
            });
            threads[i].start();
        }

        for (Thread thread : threads) {
            thread.join();
        }

        for (boolean result : results) {
            assertTrue("MAC computation should be thread-safe", result);
        }
    }
}
